#include "crypt.h"
#include "rsa_crypt.h"
#include "aes_crypt.h"

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include <error.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <poll.h>

#define SERVER_ADDR "127.0.0.1"
#define SERVER_PORT 5050

#define AES_KEY_FIXED "01234567890123456789012345678901"
#define RSA_PUBLIC_KEY_NAME "rsa_public.pem"
#define RSA_PRIVATE_KEY_NAME "rsa_private.pem"

static int sock_fd,conn_fd;
static unsigned char *rand_aes_key;
static unsigned char aes_key[AES_KEY_LEN];
static RSA *rsa_key;

#define RECV_BUF_SIZE 2048

int crypt_server(void);
int frame_parse(const unsigned char *cmd,int len);

int main(int argc,char *argv[])
{
    crypt_init();
    rsa_key = rsa_generate_key();
    if(!rsa_key)
    {
        printf("generate rsa_key failed..\r\n");
        return -1;
    }
    printf("generate rsa_key success..\r\n");
    if(rsa_save_public_to_file(rsa_key,RSA_PUBLIC_KEY_NAME) != 1)
    {
        printf("save rsa public to file failed..\r\n");
        return -1;
    }
    printf("save rsa public to file success..\r\n");
    if(rsa_save_private_to_file(rsa_key,RSA_PRIVATE_KEY_NAME,NULL) != 1)
    {
        printf("save rsa private to file failed..\r\n");
        return -1;
    }
    printf("save rsa private to file success..\r\n");
    memcpy(aes_key,AES_KEY_FIXED,strlen(AES_KEY_FIXED));
    rand_aes_key = (unsigned char *)calloc(1,AES_KEY_LEN);
    if(!rand_aes_key)
    {
        printf("calloc failed..\r\n");
        return -1;
    }
    //生成随机aes密钥
    aes_generate_rand_key(rand_aes_key);
    printf("rand_aes_key:\r\n");
    for(int i = 0;i < AES_KEY_LEN;i++)
    {
        if(i && (!(i % 10)))
        {
            printf("\r\n");
        }
        printf("%02x ",rand_aes_key[i]);
    }
    printf("\r\n");
    crypt_server();
    
    return 0;
}

int crypt_server(void)
{
    int ret = -1;
    struct sockaddr_in server_addr;
    unsigned char buf[sizeof(struct in_addr)];
    uint8_t *recv_buf;
    ssize_t len;
    struct pollfd fds[2];
    char input_buf[12];

    printf("run server:\r\n");

    recv_buf = (uint8_t *)calloc(1,RECV_BUF_SIZE);
    if(!recv_buf)
    {
        printf("calloc failed..\r\n");
        return ret;
    }

    sock_fd = socket(AF_INET,SOCK_STREAM,0);
    if(sock_fd < 0)
    {
        perror("socket ");
        return ret;
    }

    int opt = 1;
    setsockopt(sock_fd,SOL_SOCKET,SO_REUSEADDR,&opt,sizeof(opt));

    ret = inet_pton(AF_INET,SERVER_ADDR,buf);
    if(ret != 1)
    {
        printf("inet_pton error..\r\n");
        goto close_sock_fd;
    }
    server_addr.sin_addr.s_addr = ((struct in_addr *)buf)->s_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(SERVER_PORT);
    ret = bind(sock_fd,(struct sockaddr *)&server_addr,sizeof(server_addr));
    if(ret < 0)
    {
        perror("bind ");
        goto close_sock_fd;
    }
    ret = listen(sock_fd,250);
    if(ret < 0)
    {
        perror("listen ");
        goto close_sock_fd;
    }

    printf("wait client connect..\r\n");

    conn_fd = accept(sock_fd,NULL,NULL);
    if(conn_fd < 0)
    {
        perror("accept ");
        goto close_sock_fd;
    }

    printf("client connect success..\r\n");

    fds[0].fd = 0;
    fds[0].events = POLLIN;
    fds[1].fd = conn_fd;
    fds[1].events = POLLIN;
    for(;;)
    {
        ret = poll(fds,sizeof(fds) / sizeof(*fds),-1);
        if(ret < 0)
        {
            perror("poll ");
            goto close_conn_fd;
        }
        else if(!ret)
        {
            /*timeout*/
        }
        else
        {
            for(int i = 0;i < sizeof(fds) / sizeof(*fds);i++)
            {
                if(fds[i].revents & fds[i].events)
                {
                    if(fds[i].fd == 0)
                    {
                        len = read(0,input_buf,sizeof(input_buf));
                        if(len < 0)
                        {
                            if((len != EINTR) && (len != EAGAIN) && (len != EWOULDBLOCK))
                            {
                                perror("read ");
                                goto close_conn_fd;
                            }
                        }
                        else if(!len)
                        {

                        }
                        else
                        {
                            unsigned char *ciphertext;
                            ciphertext = (unsigned char *)calloc(1,len);
                            aes_encrypt((const unsigned char *)input_buf,len,aes_key,ciphertext);
                            
                            send(conn_fd,ciphertext,len,0);
                            free(ciphertext);
                        }
                    }
                    else if(fds[i].fd == conn_fd)
                    {
                        memset(recv_buf,0,RECV_BUF_SIZE);
                        len = recv(conn_fd,recv_buf,RECV_BUF_SIZE,0);
                        if(len < 0)
                        {
                            if((len != EINTR) && (len != EAGAIN) && (len != EWOULDBLOCK))
                            {
                                perror("recv ");
                                goto close_conn_fd;
                            }
                            continue;
                        }
                        else if(!len)
                        {
                            printf("client shutdown..\r\n");
                            goto close_conn_fd;
                        }
                        else
                        {
                            unsigned char *plaintext;
                            plaintext = (unsigned char *)calloc(1,len + 1);
                            aes_decrypt(recv_buf,len,aes_key,plaintext);

                            plaintext[len] = '\0';
                            printf("recv:%s\r\n",plaintext);
                            frame_parse(plaintext,len);
                            free(plaintext);
                        }
                    }
                }
            }
        }
    }

    ret = 0;
close_conn_fd:
    shutdown(conn_fd,SHUT_RDWR);
    close(conn_fd);
close_sock_fd:
    close(sock_fd);

    return ret;
}

int frame_parse(const unsigned char *cmd,int len)
{
    char *cmd_s;
    FILE *fp;
    struct stat sb;
    char *buf;
    int tmplen;
    unsigned char *ciphertext;
    unsigned char *ciphertext_rsa;
    char *prefix;
    unsigned char *sig;
    unsigned int siglen;

    cmd_s = "apply key";
    if(!strncmp((const char *)cmd,cmd_s,strlen(cmd_s)))
    {
        prefix = "rsa key";
        fp = fopen(RSA_PUBLIC_KEY_NAME,"r");
        if(!fp)
        {
            perror("fopen ");
            return -1;
        }
        fstat(fileno(fp),&sb);
        tmplen = sb.st_size;
        buf = (char *)calloc(1,tmplen + strlen(prefix));
        if(!buf)
        {
            printf("calloc failed..\r\n");
            return -1;
        }
        strncpy(buf,prefix,strlen(prefix));
        fread(buf + strlen(prefix),1,tmplen,fp);
        
        tmplen += strlen(prefix);
        ciphertext = (unsigned char *)calloc(1,tmplen);
        aes_encrypt((const unsigned char *)buf,tmplen,aes_key,ciphertext);
        
        send(conn_fd,ciphertext,tmplen,0);
        free(ciphertext);
        free(buf);
    }

    cmd_s = "public key ok";
    if(!strncmp((const char *)cmd,cmd_s,strlen(cmd_s)))
    {
        prefix = "aes key";
        
        //申请存放rsa加密密文的内存
        ciphertext_rsa = (unsigned char *)calloc(1,CIPHERTEXT_LEN(rsa_key));
        //使用rsa加密生成的aes密钥
        tmplen = rsa_private_encrypt((const unsigned char *)rand_aes_key,AES_KEY_LEN,rsa_key,ciphertext_rsa);
        //签名
        sig = (unsigned char *)calloc(1,CIPHERTEXT_LEN(rsa_key));
        if(1 != rsa_sign(rsa_key,rand_aes_key,AES_KEY_LEN,sig,&siglen))
        {
            printf("signature failed..\r\n");
            sig = NULL;
            siglen = 0;
        }
        buf = (char *)calloc(1,tmplen + strlen(prefix) + siglen);
        strncpy(buf,prefix,strlen(prefix));
        memcpy(buf + strlen(prefix),ciphertext_rsa,tmplen);
        if(siglen)
        {
            printf("siglen:%d\r\n",siglen);
            printf("signature:\r\n");
            for(int i = 0;i < siglen;i++)
            {
                if(i && (!(i % 10)))
                {
                    printf("\r\n");
                }
                printf("%02x ",sig[i]);
            }
            printf("\r\n");
            memcpy(buf + strlen(prefix) + tmplen,sig,siglen);
        }
        tmplen += strlen(prefix) + siglen;
        ciphertext = (unsigned char *)calloc(1,tmplen);
        //再将rsa加密后的密钥进行aes加密
        aes_encrypt((const unsigned char *)buf,tmplen,aes_key,ciphertext);
        //发送到客户端
        send(conn_fd,ciphertext,tmplen,0);
        free(ciphertext);
        free(buf);
        free(ciphertext_rsa);
        if(sig)
        {
            free(sig);
        }
    }

    cmd_s = "aes key ok";
    if(!strncmp((const char *)cmd,cmd_s,strlen(cmd_s)))
    {
        //更换密钥
        memcpy(aes_key,rand_aes_key,AES_KEY_LEN);
        //释放后便不能再次申请
        // free(rand_aes_key);
    }

    return 0;
}
